from elasticsearch import Elasticsearch
import os

# 初始化 Elasticsearch 客户端
def get_es_client():
    """
    获取 Elasticsearch 客户端实例
    """
    es_host = os.getenv("ES_HOST", "localhost")
    es_port = os.getenv("ES_PORT", "9200")
    es_user = os.getenv("ES_USER", None)
    es_password = os.getenv("ES_PASSWORD", None)
    
    if es_user and es_password:
        es = Elasticsearch(
            [f"http://{es_host}:{es_port}"],
            basic_auth=(es_user, es_password)
        )
    else:
        es = Elasticsearch([f"http://{es_host}:{es_port}"])
    
    return es

# 创建索引映射
def create_health_index_mapping(es_client, index_name="health_advice"):
    """
    创建健康建议索引的映射配置
    
    参数:
        es_client: Elasticsearch 客户端实例
        index_name (str): 索引名称
    """
    mapping = {
        "mappings": {
            "properties": {
                "text_content": {
                    "type": "text",
                    "analyzer": "ik_max_word",
                    "search_analyzer": "ik_smart"
                },
                "metadata": {
                    "properties": {
                        "disease": {"type": "keyword"},
                        "advice_type": {"type": "keyword"},
                        "chunk_id": {"type": "keyword"}
                    }
                }
            }
        }
    }
    
    # 创建新索引（如果不存在）
    try:
        es_client.indices.create(index=index_name, body=mapping)
    except Exception as e:
        # 索引可能已存在，忽略错误
        pass

# 存储文档到 Elasticsearch
def store_chunk_in_es(es_client, chunk_id, text_content, disease, advice_type, index_name="health_advice"):
    """
    将文本块存储到 Elasticsearch
    
    参数:
        es_client: Elasticsearch 客户端实例
        chunk_id (str): 文本块唯一标识
        text_content (str): 文本内容
        disease (str): 疾病名称
        advice_type (str): 建议类型（饮食/运动/用药/预防）
        index_name (str): 索引名称
    """
    doc = {
        "text_content": text_content,
        "metadata": {
            "disease": disease,
            "advice_type": advice_type,
            "chunk_id": chunk_id
        }
    }
    
    response = es_client.index(
        index=index_name,
        id=chunk_id,
        document=doc
    )
    
    return response

# 搜索 Elasticsearch 中的文档
def search_es_for_advice(es_client, query_text, disease=None, advice_type=None, index_name="health_advice"):
    """
    在 Elasticsearch 中搜索健康建议
    
    参数:
        es_client: Elasticsearch 客户端实例
        query_text (str): 查询文本
        disease (str): 可选，过滤特定疾病
        advice_type (str): 可选，过滤特定建议类型
        index_name (str): 索引名称
    """
    # 构建查询体
    query_body = {
        "query": {
            "bool": {
                "must": [
                    {
                        "match": {
                            "text_content": query_text
                        }
                    }
                ]
            }
        }
    }
    
    # 添加过滤条件
    filters = []
    if disease:
        filters.append({"term": {"metadata.disease": disease}})
    if advice_type:
        filters.append({"term": {"metadata.advice_type": advice_type}})
    
    if filters:
        query_body["query"]["bool"]["filter"] = filters
    
    # 执行搜索
    response = es_client.search(
        index=index_name,
        body=query_body
    )
    
    return response